#!/bin/bash
# vim: set et sw=4 sts=4 tw=80:
# Copyright 2007-2009 Gentoo Foundation
# Distributed under the terms of the GNU General Public License v2

# A bit of hackery to update everything that is humanly possible
# that maybe related to an older version of python. This script can
# be run as many times as you like.
#
# OLD_PY_VER      = old python version we are upgrading from
# NEW_PY_VER      = new python version we are upgrading to
# PKGS_EXCEPTIONS = packages that should NOT be re-emerged for any reason
# PKGS_MANUAL     = packages that should be re-emerged even if they don't
#                   fit the criteria (eg. ones that have python compiled
#                   statically)
#
# Runtime Variables:
#
# PKGS_TO_REMERGE = list of packages we deem to need re-emerging
# PKGS_OK         = list of packages that should be merged without any problems
# PKGS_MISSING    = list of packages that are installed, but cannot be merged
#                   because they have been pruned from portage
# PKGS_MASKED     = list of packages that are installed, but masked.
#

VERSION="0.7"
NEW_PY_VER=$(/usr/bin/python -V 2>&1 | sed 's:Python ::' | cut -d. -f1-2)

PKGS_EXCEPTIONS="dev-lang/python sys-apps/portage"
PKGS_MANUAL="app-office/gnumeric app-office/dia dev-libs/boost x11-libs/vte"

PRETEND=0
REINSTALL_IDENTICAL_VERSIONS=0
VERBOSE=0
PKGS_TO_REMERGE=""
PKGS_COUNT_REMERGE=0
PORTAGE_PYTHON="/usr/bin/python"
PYTHON_VERSIONS="3.2 3.1 3.0 2.7 2.6 2.5 2.4 2.3 2.2 2.1"

SUPPORTED_PMS="portage pkgcore paludis"
PMS_COMMAND=( "emerge" "pmerge" "paludis" )
PMS_OPTIONS=( "-vD1" "-Do" "-i1" )
PMS_INDEX=0
CUSTOM_PMS_COMMAND=""
ADDITIONAL_OPTIONS=""

# Checks
CHECK_ECLASS=0
CHECK_MANUAL=1
CHECK_PYLIBDIR=1
CHECK_SONAME=1
CHECK_ECLASS_NEED_REBUILD=1

# load the gentoo-style info macros, but hack to get around
# it thinking this is an rc script
EBUILD="1"
source /etc/init.d/functions.sh

# portage variables
PKG_DBDIR=/var/db/pkg

# usage()
# display usage
usage() {
    cat <<EOF_USAGE
${0##*/} -- Find & rebuild packages broken due to a python upgrade

Usage: python-updater [OPTION]

Options:
    -h, --help      Print usage
    -V, --version   Print version
    -p, --pretend   Pretend (don't do anything)
    -v, --verbose   Increase verbosity (may be specified multiple times)
    -o PYVER, --old-version PYVER
                    Set old python version to upgrade from to PYVER
    --reinstall-identical-versions
                    Reinstall identical versions of packages
    -P PM, --package-manager PM
                    Use package manager PM, where PM can be one of:
$(for p in ${SUPPORTED_PMS}; do
echo -ne $'\t\t    '\* ${p}
if [[ ${p} == portage ]]; then
    echo ' (Default)'
else
    echo
fi
done)
    -c CMD, --command CMD
                    Pipe found packages to command CMD instead of invoking package
                    manager. Only for debug and script use.
    --package-manager-command CMD
                    Call CMD instead of the default command for the specified
                    package manager.
    -eCHECK --enable-CHECK
                    Enable CHECK where CHECK can be one of:
                    * eclass       (Disabled by default)
                    * pylibdir     (Enabled by default)
                    * soname       (Enabled by default)
                    * manual       (Enabled by default)
                    * need_rebuild (Enabled by default)
    -dCHECK --disable-CHECK
                    Disable CHECK. See --enable option for a list of checks
    -- OPTIONS      Pass additional options to package manager

See CHECKS section in the manpage for explanations about checks and
EXAMPLES section to learn how to use them.
EOF_USAGE
}

# veinfo(verbosity, message)
# einfo message if VERBOSE is bigger than verbosity
veinfo() {
    if [[ VERBOSE -ge $1 ]]; then
        shift
        einfo $@
    fi
}

# get_old_pyver()
# Find old python version, return non-zero if not found
get_old_pyver() {
    for old in ${PYTHON_VERSIONS}; do
        if [[ "${old}" != "${NEW_PY_VER}" && -e /usr/bin/python${old} ]]
        then
                echo -n "${old}"
                return 0
        fi
    done
    eerror "Couldn't determine any previous Python version(s)."
    eerror "Use -o OLD_PYTHON_VERSION to specify your old python version."
    return 1
}

# get_portage_python(oldpy,newpy)
# Find where portage is, in pythonX.Y or somewhere else?
get_portage_python() {
    if [[ $# -lt 2 ]]; then
        eerror "get_portage_python(): Missing arguments."
        exit 1
    fi
    local oldpy="$1" newpy="$2"

    pybin=/usr/bin/python
    for py in ${pybin} ${pybin}${oldpy} ${pybin}${newpy}; do
        if ${py} -c "import portage" > /dev/null 2>&1; then
            echo -n "${py}"
            return 0
        fi
    done
    eerror "Couldn't determine portage python"
    return 1
}

# get_portage_portdir()
# Check if portage knows about PORTDIR and return it
get_portage_portdir() {
    local portdir="$(/usr/bin/portageq portdir)"

    if [[ -z "${portdir}" ]]; then
        eerror "Unable to proceed. Can not find PORTDIR. Make sure the command:"
        eerror " "
        eerror "  portageq portdir"
        eerror "returns a value. If it doesn't, make sure you have updated to"
        eerror "latest portage version."
        eerror " "
        eerror "Report bugs to http://bugs.gentoo.org/"
        return 1
    else
        echo -n "${portdir}"
        return 0
    fi
}

# Respect PYUPDATER_OPTIONS
if [[ -n "${PYUPDATER_OPTIONS}" ]]; then
    set -- ${PYUPDATER_OPTIONS} $@
fi

# Respect PACKAGE_MANAGER environment variable
i=0
for pm in ${SUPPORTED_PMS}; do
    if [[ "${pm}" == "${PACKAGE_MANAGER}" ]]; then
        PMS_INDEX="${i}"
        break
    fi
    ((i++))
done

# Command Line Parsing
while [[ -n "$1" ]]; do
    case "$1" in
        -h|--help)
            usage
            exit 0
            ;;
        -V|--version)
            echo "${VERSION}"
            exit 0
            ;;
        -p|--pretend)
            PRETEND=1
            ;;
        -v|--verbose)
            VERBOSE=$(( $VERBOSE + 1 ))
            ;;
        -o|--old-version)
            shift
            OLD_PY_VER="$1"
            ;;
        --reinstall-identical-versions)
            REINSTALL_IDENTICAL_VERSIONS=1
            ;;
        -P|--package-manager)
            shift
            PACKAGE_MANAGER="$1"
            case "${PACKAGE_MANAGER}" in
                portage|pkgcore|paludis)
                    ;;
                *)
                    echo "unrecognised package manager selected. please select between ${SUPPORTED_PMS}"
                    exit
                    ;;
            esac

            # PMS_INDEX is used to select the right commands and options for the selected package manager
            PMS_INDEX=0
            for PM in ${SUPPORTED_PMS}; do
                [[ ${PM} == ${PACKAGE_MANAGER} ]] && break
                PMS_INDEX=$((${PMS_INDEX} + 1))
            done
            ;;
        --package-manager-command)
            shift
            CUSTOM_PMS_COMMAND="$1"
            ;;
        -c|--command)
            shift
            PIPE_COMMAND="$1"
            ;;
        -ee*|--enable-e*)
            CHECK_ECLASS=1
            ;;
        -de*|--disable-e*)
            CHECK_ECLASS=0
            ;;
        -em*|--enable-m*)
            CHECK_MANUAL=1
            ;;
        -dm*|--disable-m*)
            CHECK_MANUAL=0
            ;;
        -ep*|--enable-p*)
            CHECK_PYLIBDIR=1
            ;;
        -dp*|--disable-p*)
            CHECK_PYLIBDIR=0
            ;;
        -es*|--enable-s*)
            CHECK_SONAME=1
            ;;
        -ds*|--disable-s*)
            CHECK_SONAME=0
            ;;
        -en*|--enable-n*)
            CHECK_ECLASS_NEED_REBUILD=1
            ;;
        -dn*|--disable-n*)
            CHECK_ECLASS_NEED_REBUILD=0
            ;;
        --)
            shift
            ADDITIONAL_OPTIONS="$*"
            break
            ;;
        *)
            usage
            echo "unrecognised option: $1"
            exit 0
            ;;
    esac
    shift
done

# Sanity check
PORTDIR="$(get_portage_portdir)"
[[ $? != 0 ]] && exit 1

# Determine old python version
if [[ -z "${OLD_PY_VER}" ]]; then
    OLD_PY_VER="$(get_old_pyver)"
    if [[ $? != 0 ]]; then
        exit 1
    fi
fi

# Get portage python
PORTAGE_PYTHON="$(get_portage_python ${OLD_PY_VER} ${NEW_PY_VER})"
[[ $? != 0 ]] && exit 1


einfo "Starting Python Updater from ${OLD_PY_VER} to ${NEW_PY_VER} :"
if [[ CHECK_SONAME -ne 0 ]]; then
    if ! type -P scanelf >/dev/null 2>&1; then
        ewarn "scanelf not found!"
        ewarn "check soname is disabled."
        CHECK_SONAME=0
    else
        veinfo 1 'check "soname" enabled.'
        OLD_SONAME="$(readlink -n /usr/lib/libpython${OLD_PY_VER}.so)"
        if [[ -z "${OLD_SONAME}" ]]; then
            ewarn "Couldn't find old libpython soname"
            ewarn "Disabling soname check."
            CHECK_SONAME=0
        fi
    fi
else
    veinfo 1 'check "soname" disabled.'
fi
[[ CHECK_PYLIBDIR -ne 0 ]] \
    && veinfo 1 'check "pylibdir" enabled.' \
    || veinfo 1 'check "pylibdir" disabled.'
[[ CHECK_ECLASS -ne 0 ]] \
    && veinfo 1 'check "eclass" enabled.' \
    || veinfo 1 'check "eclass" disabled.'
[[ CHECK_MANUAL -ne 0 ]] \
    && veinfo 1 'check "manual" enabled.' \
    || veinfo 1 'check "manual" disabled.'
[[ CHECK_ECLASS_NEED_REBUILD -ne 0 ]] \
    && veinfo 1 'check "need_rebuild" enabled.' \
    || veinfo 1 'check "need_rebuild" disabled.'

# Iterate through the contents of all the installed packages.
# ${PKG_DBDIR} must be followed by '/' to avoid problems when ${PKG_DBDIR} is a symlink.
for content in `find ${PKG_DBDIR}/ -name CONTENTS`; do
    # extract the category, package name and package version
    CATPKGVER=$(echo ${content} | sed "s:${PKG_DBDIR}/\(.*\)/CONTENTS:\1:")
    CATPKG="${CATPKGVER%%-[0-9]*}"
    veinfo 2 "Checking ${CATPKGVER}"

    # exclude packages that are an exception, like portage and python itself.
    exception=0
    for exp in ${PKGS_EXCEPTIONS}; do
        if [[ -z "${CATPKG##${exp}}" ]]; then
            veinfo 2 "Skipping ${CATPKG}, reason: exception"
            exception=1
            break;
        fi
    done

    [[ ${exception} == 1 ]] && continue

    # Check if package is in PKGS_MANUAL
    if [[ CHECK_MANUAL -ne 0 ]]; then
        for pkg in ${PKGS_MANUAL}; do
            if [[ -z "${CATPKG##${pkg}}" ]]; then
                exception=2
                break;
            fi
        done
    fi

    # Replace SLOT by version number when REINSTALL_IDENTICAL_VERSIONS == 1
    # Reinstall identical versions when SLOT doesn't exist, bug #201848
    if [[ REINSTALL_IDENTICAL_VERSIONS -eq 1 || ! -f "${content/CONTENTS/SLOT}" ]]; then
        CATPKGVER="=${CATPKGVER}"
    else
        SLOT=$(< ${content/CONTENTS/SLOT})
        CATPKGVER="${CATPKG}:${SLOT}"
    fi

    if [[ ${exception} = 2 ]]; then
        PKGS_TO_REMERGE="${PKGS_TO_REMERGE} ${CATPKGVER}"
        eindent
        einfo "Adding to list: ${CATPKGVER}"
        eindent
        einfo "check: manual [Added to list manually, see CHECKS in manpage for more information.]"
        eoutdent && eoutdent
        continue
    fi

    if [[ CHECK_ECLASS_NEED_REBUILD -ne 0 ]]; then
        ENVIRON="${content/CONTENTS/environment.bz2}"
        if bzip2 -dc ${ENVIRON} | grep -qe "^\(export \)\?PYTHON_NEED_REBUILD=${OLD_PY_VER}"; then
            PKGS_TO_REMERGE="${PKGS_TO_REMERGE} ${CATPKGVER}"
            eindent
            einfo "Adding to list: ${CATPKGVER}"
            eindent
            veinfo 1 "check: need_rebuild [ Ebuild set PYTHON_NEED_REBUILD=${OLD_PY_VER} ]"
            eoutdent && eoutdent
            continue
        fi
    fi

    if [[ CHECK_PYLIBDIR -ne 0 ]]; then
        # Search for possible old python dirs in CONTENTS
        # /usr/include/python$old
        # /usr/lib/python$old
        # /usr/lib32/python$old
        # /usr/lib64/python$old
        if grep -qe "/usr/\(include\|lib\(32\|64\)\?\)/python${OLD_PY_VER}" ${content}; then
            PKGS_TO_REMERGE="${PKGS_TO_REMERGE} ${CATPKGVER}"
            eindent
            einfo "Adding to list: ${CATPKGVER}"
            eindent
            veinfo 1 "check: pylibdir [ Installed file under old python library directory ]"
            eoutdent && eoutdent
            continue
        fi
    fi

    if [[ CHECK_SONAME -ne 0 ]]; then
        broken_libs="$(scanelf -qBN ${OLD_SONAME} < <(
            grep -e '^obj' ${content} | cut -d' ' -f2))"
        if [[ -n "${broken_libs}" ]]; then
            PKGS_TO_REMERGE="${PKGS_TO_REMERGE} ${CATPKGVER}"
            eindent
            einfo "Adding to list: ${CATPKGVER}"
            eindent
            veinfo 1 "check: soname [ Libraries linked to old libpython found:"
            veinfo 1 "${broken_libs}"
            veinfo 1 "]"
            eoutdent && eoutdent
        fi
    fi

    if [[ CHECK_ECLASS -ne 0 ]]; then
        ENVIRON="${content/CONTENTS/environment.bz2}"
        if bzip2 -dc ${ENVIRON} | grep -qe "^\(export \)\?PYVER=${OLD_PY_VER}"; then 
            PKGS_TO_REMERGE="${PKGS_TO_REMERGE} ${CATPKGVER}"
            eindent
            einfo "Adding to list: ${CATPKGVER}"
            eindent
            veinfo 1 "check: eclass [ Ebuild set PYVER=${OLD_PY_VER} ]"
            eoutdent && eoutdent
            continue
        fi
    fi
done

# Pipe to command if we have one
if [[ -n "${PIPE_COMMAND}" ]]; then
    echo "${PKGS_TO_REMERGE}" | ${PIPE_COMMAND}
    exit $?
fi

# Filter out --getbinpkg, --getbinpkgonly, --usepkg and --usepkgonly options in EMERGE_DEFAULT_OPTS environment variable
emerge_default_opts=""
for option in ${EMERGE_DEFAULT_OPTS}; do
    if [[ "${option}" == -[[:alnum:]]* ]]; then
        emerge_default_opts+=" ${option//[gGkK]/}"
    elif [[ "${option}" != "--getbinpkg" && "${option}" != "--getbinpkgonly" && "${option}" != "--usepkg" && "${option}" != "--usepkgonly" ]]; then
        emerge_default_opts+=" ${option}"
    fi
done
export EMERGE_DEFAULT_OPTS="${emerge_default_opts# }"

# only pretending?
[[ PRETEND -eq 1 ]] && PMS_OPTIONS[${PMS_INDEX}]="${PMS_OPTIONS[${PMS_INDEX}]}p"

# (Pretend to) remerge packages
if [[ -n "${PKGS_TO_REMERGE}" ]]; then
    pmscmd="${CUSTOM_PMS_COMMAND}"
    [[ -z "${pmscmd}" ]] && pmscmd="${PMS_COMMAND[${PMS_INDEX}]}"
    cmd="${pmscmd} ${PMS_OPTIONS[${PMS_INDEX}]} ${PKGS_TO_REMERGE} ${ADDITIONAL_OPTIONS}"
    einfo ${cmd}
    ${cmd}
else
    einfo "No packages needs to be remerged."
fi

